/***************************************************************************
 *  Title: Kernel Memory Allocator
 * -------------------------------------------------------------------------
 *    Purpose: Kernel memory allocator based on the buddy algorithm
 *    Author: Stefan Birrer
 *    Copyright: 2004 Northwestern University
 ***************************************************************************/
/***************************************************************************
 *  ChangeLog:
 * -------------------------------------------------------------------------
 *    Revision 1.2  2009/10/31 21:28:52  jot836
 *    This is the current version of KMA project 3.
 *    It includes:
 *    - the most up-to-date handout (F'09)
 *    - updated skeleton including
 *        file-driven test harness,
 *        trace generator script,
 *        support for evaluating efficiency of algorithm (wasted memory),
 *        gnuplot support for plotting allocation and waste,
 *        set of traces for all students to use (including a makefile and README of the settings),
 *    - different version of the testsuite for use on the submission site, including:
 *        scoreboard Python scripts, which posts the top 5 scores on the course webpage
 *
 *    Revision 1.1  2005/10/24 16:07:09  sbirrer
 *    - skeleton
 *
 *    Revision 1.2  2004/11/05 15:45:56  sbirrer
 *    - added size as a parameter to kma_free
 *
 *    Revision 1.1  2004/11/03 23:04:03  sbirrer
 *    - initial version for the kernel memory allocator project
 *
 ***************************************************************************/

#ifdef KMA_BUD
#define __KMA_IMPL__

/************System include***********************************************/

#include <assert.h>
#include <stdlib.h>
#include <stdio.h>

/************Private include**********************************************/

#include "kma_page.h"
#include "kma.h"

/************Defines and Typedefs*****************************************/
/*  #defines and typedefs should have their names in all caps.
 *  Global variables begin with g. Global constants with k. Local
 *  variables should be in all lower case. When initializing
 *  structures and arrays, line everything up in neat columns.
 */

/************Global Variables*********************************************/

/************Function Prototypes******************************************/
	
/************External Declaration*****************************************/

/**************Implementation***********************************************/

/*  FOR EACH PAGE:
 * Smallest trunk = 64 bytes. We need to store a bitmap, a pointer to the page
 * struct, two pointers to the right and left child at the beginning of the
 * page. If we use char to represent bool, we have 8192 - 8192/64 - 8 =
 * 8056 bytes left. I have checked that the max request is about 8000 bytes.
 * The remaining space should be enough to  satisfy the max request. I am
 * using a BST to store the pages using page id as key.
 *
 * FOR THE EACH MEMORY SEGMENT WITHIN A PAGE:
 * I use an array for the freelist headers and a binary search tree for the
 * buffer list. At the beginning of each segment, a number x (key) =
 * location in page * page number, a pointer to right child and a pointer to
 * left child is stored. There are functions to convert x to the page number
 * and its corresponding position in the page. The list is sorted according
 * to this number x.
 */

// it is an integer list because it stores the addr of the first poistion of each memory segment
// which is an integer
int *freeList[7] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL}; // 64, 128, 256, 512, 1024, 2048, 4096

////////////////////CODE BELOW FOR HANDLING TREE OPERATIION///////////////////////

// addr = address of the memory segment
// NULL should already be set at addr + 1 and addr + 2
void BSTInsertFreeSeg(int key, int **root, int *addr) {
	// the memory looks like |int|pt to left child|pt to right child|memory allocated to user|
	// p points to the first position
	int *p = *root;
	int *pp = 0; // parent of p

	while (p) { // p != NULL
		pp = p;
		if (key < *p)
			p = *((int**)(p + 1)); // size of pointer and size of int same in 32 bit sys
		else if (key > *p)
			p = *((int**)(p + 2));
	}

	if (*root) { // if the head is not NULL
		if (key < *pp)
			*((int**)(pp + 1)) = addr; // store as the addr of right child of parent
		else
			*((int**)(pp + 2)) = addr; // store as the addr of left child of parent
	}
	else
		*root = addr;
}

// when memory is allocated from the buffer, delete and return the location with the largest addr
// so that we can allocate the left segment of a page first to prevent problems caused by bitmap
// return the "virtual address" so that we can convert it to the real address
int* BSTDeleteFreeSeg(int** root) {
	int *p = *root;
	int *pp = 0;
	int *memAllocated = 0;

	// find the piece with the largest virtual addr
	while (*((int**)(p + 2))) {
		pp = p;
		p = *((int**)(p + 2));
	}

	memAllocated = p;

	// remove the piece from our tree
	int *c;
	c = *((int**)(p + 1)); // p has left child only, c = right child of p

	if (p == *root)
		*root = c;
	else
		*((int**)(pp + 2)) = c; // p must be its parent's right child as it is max

	return memAllocated;
}

// given the pageId, delete all the segment in the tree
void BSTDeleteAllFreeSegOfPage(int** root, int pageId) {
	int low = 0 + 8192 * pageId;
	int high = 8191 + 8192 * pageId;

	while (TRUE) {
		int *p = *root;
		int *pp = 0;
		while (p && (*p > high || *p < low)) {
			pp = p;
			if (high < *p)
				p = *((int**)(p + 1));
			else if (low > *p)
				p = *((int**)(p + 2));
		}

		if (!p) // cannot  be found
			return;

		if (*((int**)(p + 1)) && *((int**)(p + 2))) {
			int *s = *((int**)(p + 1));
			int *ps = p;
			while (*((int**)(s + 2))) {
				ps = s;
				s = *((int**)(s + 2));
			}

			if (pp != NULL) {
				if (p == *((int**)(pp + 1)))
					*((int**)(pp + 1)) = s;
				else
					*((int**)(pp + 2)) = s;
			}

			if (ps != p) {
				if (s == *((int**)(ps + 1)))
					*((int**)(ps + 1)) = p;
				else
					*((int**)(ps + 2)) = p;
			}

			int *pLeft = *((int**)(p + 1));
			int *pRight = *((int**)(p + 2));

			*((int**)(p + 1)) = *((int**)(s + 1));
			*((int**)(p + 2)) = *((int**)(s + 2));

			if (pLeft != s)
				*((int**)(s + 1)) = pLeft;
			else
				*((int**)(s + 1)) = p;
			if (pRight != s)
				*((int**)(s + 2)) = pRight;
			else
				*((int**)(s + 2)) = p;

			if (ps == p)
				pp = s;
			else
				pp = ps;
		}

		int *c = 0;
		if (*((int**)(p + 1)))
			c = *((int**)(p + 1));
		else
			c = *((int**)(p + 2));

		if (p == *root)
			*root = c;
		else {
			if (p == *((int**)(pp + 1)))
				*((int**)(pp + 1)) = c;
			else
				*((int**)(pp + 2)) = c;
		}
	}
}

// try to search and delete the buddy
// if deleted return true else return false
// key is the virtual addr = page id * addr in the page
bool canDeleteBuddy(int** root, int key) {
	int *p = *root;
	int *pp = 0;

	while (p && *p != key) { // the pointer is not null and the number the pointer pointing to not equals key
		pp = p;
		if (key < *p)
			p = *((int**)(p + 1)); // key less then the number the pt pting to, go to left
		else
			p = *((int**)(p + 2));
	}

	if (!p) // the pointer is null, we cannot fine the node
		return FALSE;

	if (*((int**)(p + 1)) && *((int**)(p + 2))) { // have both left and right child
		int *s = *((int**)(p + 1)); // s points to the left child
		int *ps = p; // ps is the parent of s
		while (*((int**)(s + 2))) { // the right child of s is not null
			ps = s;
			s = *((int**)(s + 2));
		}

		if (pp != NULL) {
			if (p == *((int**)(pp + 1)))
				*((int**)(pp + 1)) = s;
			else
				*((int**)(pp + 2)) = s;
		}
		if (ps != p) {
			if (s == *((int**)(ps + 1)))
				*((int**)(ps + 1)) = p;
			else
				*((int**)(ps + 2)) = p;
		}

		int *pLeft = *((int**)(p + 1));
		int *pRight = *((int**)(p + 2));

		*((int**)(p + 1)) = *((int**)(s + 1));
		*((int**)(p + 2)) = *((int**)(s + 2));

		if (pLeft != s)
			*((int**)(s + 1)) = pLeft;
		else
			*((int**)(s + 1)) = p;
		if (pRight != s)
			*((int**)(s + 2)) = pRight;
		else
			*((int**)(s + 2)) = p;

		if (ps == p)
			pp = s;
		else
			pp = ps;
	}

	int *c = 0;
	if (*((int**)(p + 1)))
		c = *((int**)(p + 1));
	else
		c = *((int**)(p + 2));

	if (p == *root)
		*root = c;
	else {
		if (p == *((int**)(pp + 1)))
			*((int**)(pp + 1)) = c;
		else
			*((int**)(pp + 2)) = c;
	}

	return TRUE;
}

////////////////////CODE ABOVE FOR HANDLING TREE OPERATIION///////////////////////

// for example, convert 256 to 2, 512 to 3
// for requesting memory
int convertBufferSize2Index(int bufferSize) {
	switch (bufferSize) {
		case 64:
			return 0;
		case 128:
			return 1;
		case 256:
			return 2;
		case 512:
			return 3;
		case 1024:
			return 4;
		case 2048:
			return 5;
		case 4096:
			return 6;
		default:
			return -1;
	}
}

// for example, 1024 * 5 (position 1024 in page id = 5) to 5
// useful when all the memory in the page is not used
// we can get back the page and free it
int convertVirtualAddr2PageId(int virtualAddr) {
	int id = 0;
	while (!(virtualAddr >= 0 + 8192 * id && virtualAddr < 8191 + 8192 * id))
		id++;
	return id;
}

bool isFirstSegment(int virtualAddr) {
	int id = convertVirtualAddr2PageId(virtualAddr);
	virtualAddr = virtualAddr - 8192 * id;
	if (virtualAddr == 0)
		return TRUE;
	return FALSE;
}

int convertSize2Power2(int size) {
	int i = 0;
	while ((64 << i) < size + 4) // 4 bytes is used to store the virtual addr of the memory segment
		i++;
	return (64 << i);
}

// ptr points to start of bitmap
// starting addr is the addr of the memory segment in that page (i.e. 0- 8191)
// allocatedSize is the true allocated size
// bit = 0 -> free; bit = 1 -> allocated
void setBitmap(char *ptr, int startingAddr, int allocatedSize, char bit) {
	ptr += startingAddr / 64;
	int i;
	int numOfBits = allocatedSize / 64;

	for (i = 0; i < numOfBits; i++, ptr++)
		*ptr = bit;
}

// return addr of the memory allocated
// 3 cases:
// if not enough memory in page, allocate page, split, put to freelists
// if enough but not get from its own bucket, split what you get, put, split, put... return
// if get immediate, return
void* kma_malloc(kma_size_t size) { // kma_size_t (int) define in kma.h
	if (size >= 8192 - 128 - 12 || size <= 0) // 128 is the size of the bitmap, 12 bytes store the page id and pointers
		return NULL;

	int sizeRequired = convertSize2Power2(size); // get the next power 2
	int index = convertBufferSize2Index(sizeRequired); // get the index of bucket to find the memory seg required

	int foundIndex = index;
	if (index != -1) {
		while (freeList[foundIndex] == NULL && foundIndex <= 6)
			foundIndex++;
	}
	else
		foundIndex = 7; // we need to allocate a whole page for the memory requested; index = -1 only when we need 8192

	if (foundIndex == 7) { // cannot find non empty bucket with enough size
		kma_page_t* newPage = get_page();

		// 1. initialize the bitmap
		char *bmPtr = newPage->ptr; //pt to an 1-byte position
		int i;
		for (i = 0; i < 128; i++, bmPtr++) // set the next 128 byte to be zero (free)
			*bmPtr = 0;

		// 2. put the page to the tree
		kma_page_t** pPtr = (kma_page_t**)bmPtr;
		*pPtr = newPage; // newPage is a pointer to page

		if (sizeRequired == 8192) {// it must use the whole page
			setBitmap(newPage->ptr, 0, 8192, 1);

			return (void*)(pPtr + 3);
		}

		// 3. split the page and put the segments to the tree
		// it is the first segment of size 256, we should never touch it
		int *mPtr, **mPtr2;
		int segSize = 256;

		// split in advance
		while (segSize <= 4096) {
			mPtr = newPage->ptr; // move it to the beginning of the page
			mPtr += segSize / sizeof(int); // cut at 256, 512, 1024, 2048, 4069
			*mPtr = segSize + 8192 * newPage->id;
			mPtr2 = (int**)mPtr;
			*(mPtr2 + 1) = NULL;
			*(mPtr2 + 2) = NULL;
			BSTInsertFreeSeg(*mPtr, &freeList[convertBufferSize2Index(segSize)], mPtr);
			segSize *= 2;
		}

		foundIndex = index;
		while (freeList[foundIndex] == NULL && foundIndex <= 6)
			foundIndex++;
	}

	if (foundIndex == index) { // if we can get exactly the size that we want, no splitting is required
		int *ptr = BSTDeleteFreeSeg(&freeList[foundIndex]); // pt to the memSeg and remove it from the tree

		int virtualAddr = *ptr;
		int id = convertVirtualAddr2PageId(virtualAddr);
		int virtualAddrInPage = virtualAddr - 8192 * id;
		char* bitmapPtr = (char*)(ptr - virtualAddrInPage / sizeof(int)); // point to first position of bitmap
		setBitmap(bitmapPtr, virtualAddrInPage, sizeRequired, 1);

		ptr++; // skip 4 bytes which store the virtual addr
		return (void*)ptr;
	}

	if (foundIndex > index && foundIndex <= 6) { // splitting is required
		int *ptr = BSTDeleteFreeSeg(&freeList[foundIndex]);
		int startingAddr = *ptr;
		int currentSegSize = (64 << foundIndex); // size of memory in the bucket
		int currentBucket = foundIndex;

		while (currentSegSize != sizeRequired) { // splitting is required
			*((int**)(ptr + 1)) = NULL;
			*((int**)(ptr + 2)) = NULL;
			currentBucket--;
			BSTInsertFreeSeg(*ptr, &freeList[currentBucket], ptr);

			ptr += (64 << currentBucket) / sizeof(int); // e.g. size of bucket 1024, then divide by 2 = 512
			startingAddr = startingAddr + (64 << currentBucket);
			*ptr = startingAddr;
			*((int**)(ptr + 1)) = NULL;
			*((int**)(ptr + 2)) = NULL;

			currentSegSize = currentSegSize / 2;
		}

		int virtualAddr = *ptr;
		int id = convertVirtualAddr2PageId(virtualAddr);
		int virtualAddrInPage = virtualAddr - 8192 * id;
		char* bitmapPtr = (char*)(ptr - virtualAddrInPage / sizeof(int)); // point to first position of bitmap
		setBitmap(bitmapPtr, virtualAddrInPage, sizeRequired, 1);

		ptr++; // move by 4 bytes as it stores the virtual addr
		return (void*)ptr;
	}

	return NULL;
}
// virtualaddr is the addr in the page
bool isLeftBuddy(int low, int high, int virtualAddr, int size) {
	int mid = low + (high - low + 1) / 2;
	int segmentSize = mid - low;
	if (segmentSize == size && virtualAddr >= low && virtualAddr <= mid - 1)
		return TRUE;
	else if (segmentSize == size && virtualAddr >= mid && virtualAddr <= high)
		return FALSE;
	else if (virtualAddr < low || virtualAddr > high)
		return FALSE;

	// below means size are not the same, split again
	// or it means the virtual Addr is not within low or high
	return (isLeftBuddy(low, mid - 1, virtualAddr, size) || isLeftBuddy(mid, high, virtualAddr, size));
}

bool isAllFree(char* bitmapPtr) {
	int i;
	for (i = 0; i < 128; i++, bitmapPtr++) {
		if (*bitmapPtr == 1)
			return FALSE;
	}
	return TRUE;
}

// cannot remove the virtual addr, you need the virtual addr to check its position and to which page it belongs!!
// assume we can really deallocate the memory first
void kma_free(void* ptr, kma_size_t size) {
	int deletingSize = convertSize2Power2(size);
	if (deletingSize == 8192) { // the segment of 8192 bytes is not store in any buckets
		kma_page_t* deletingPage = *((kma_page_t**)ptr - 3);
		free_page(deletingPage);
		return;
	}

	int *intPtr = (int*)(ptr);
	int virtualAddr = *(intPtr - 1);
	int id = convertVirtualAddr2PageId(virtualAddr);
	int index = convertBufferSize2Index(deletingSize);
	int virtualAddrInPage = virtualAddr - 8192 * id;

	// set the bitmap
	char* bitmapPtr = (char*)(intPtr - 1 - virtualAddrInPage / sizeof(int)); // -1 because we use 4 bytes to store the virtual addr
	setBitmap(bitmapPtr, virtualAddrInPage, deletingSize, 0);

	// check the bitmap
	// if all are 0, remove page and the all the memory segment with that addr
	if (isAllFree(bitmapPtr)) {
		int i;
		for (i = 0; i < 7; i++)
			BSTDeleteAllFreeSegOfPage(&freeList[i], id);

		kma_page_t* deletingPage = *((kma_page_t**)(bitmapPtr + 128));
		free_page(deletingPage);

		return;
	}

	if (deletingSize != 8192) {
		// ptr should be pointing to the position after the virtaul addr (pointer to left child)
		*((int**)intPtr) = NULL;
		*((int**)(intPtr + 1)) = NULL;

		intPtr--; // points back to the virtualAddr
		bool canDelete = TRUE;
		int buddyAddr;

		while (deletingSize <= 4096 && canDelete) {
			bool isLeft = isLeftBuddy(0, 8191, virtualAddrInPage, deletingSize);

			if (isLeft) // left
				buddyAddr = virtualAddr + deletingSize;
			else // right
				buddyAddr = virtualAddr - deletingSize;

			canDelete = canDeleteBuddy(&freeList[index], buddyAddr);

			if (canDelete) {
				if (!isLeft) {
					virtualAddrInPage -= deletingSize;
					virtualAddr = buddyAddr;
					intPtr = intPtr - deletingSize / sizeof(int);
				}
				deletingSize *= 2;
				index++;
			}
		}
		*intPtr = virtualAddr;
		*((int**)(intPtr + 1)) = NULL;
		*((int**)(intPtr + 2)) = NULL;
		BSTInsertFreeSeg(virtualAddr, &freeList[index], intPtr); // insert
	}

	return;
}

#endif // KMA_BUD
